open Curses

exception Of_representation of string

module Specification = 
struct
  type t' = 
    | Variant of (string * (t list)) list
    | Record of (string * t) list
    | Tuple of t list
    | Atom of (string -> string list)
  and t = unit -> t'
end

module Representation = 
struct
  type t = 
    | Variant of string * t list
    | Record of (string * t) list
    | Tuple of t list
    | Atom of string
    | Hole
end

exception Unexpected
exception Hole 


module Implementation = 
struct
  type choices = 
    | Variant of string * node list
    | Record of (string * node) list
    | Tuple of node list
    | Atom of string
  and node = 
      {mutable body: choices option; 
       mutable str: string; 
       id:int;
       possibilities:(string -> ((unit -> choices) * string option) list);
       mutable possibility_list: (string * (unit -> choices)) list
      }
	
  type t = unit -> node

  type possibility = string -> ((unit -> choices) * string option) list
  type possibility_result = ((unit -> choices) * string option) list


  let counter = ref 0 
  ;;
  
  
  let get_uid () = 
    let ret = !counter in
      counter := !counter + 1;
      ret
  ;;

  let is_prefix_ic str ~prefix = 
    if String.length str < String.length prefix then false else
      let rec iter i = 
	if i = String.length prefix then true 
	else if Char.lowercase str.[i] = Char.lowercase prefix.[i] then 
	  iter (i+1)
	else
	  false
      in
	iter 0
  ;;

  let filter_map f lst = 
    let rec loop =
      function
	| [] -> 
	    []
	| x :: xs ->
	    begin match f x with
	      | None -> loop xs
	      | Some y -> y :: loop xs end
    in
      loop lst
  ;;

  
  let rec of_specification (spec:Specification.t) : t = 
    let module S = Specification in 
    let rec iter s = iter' (s ())
    and iter' : Specification.t' -> node = 
      function
	| S.Variant lst ->
	    let f str (cnstr,t_lst) : ((unit -> choices) * string option) option = 
	      let nonempty = String.length str <> 0 in 
		if nonempty && is_prefix_ic ~prefix:str cnstr then 
		  Some 
		    ((fun () -> Variant (cnstr,List.map iter t_lst)),
		     Some cnstr)
		else
		  None
	    in
	    let possibilities str = 
	      filter_map
		(f str)
		lst
	    in
	      {str="";
	       id = get_uid ();
	       body = None;
	       possibilities = possibilities;
	       possibility_list = [];
	      }
	| S.Tuple lst ->
	    let v () = Tuple (List.map (iter: S.t -> node) lst) in 
	    let possibilities str = 
	      if is_prefix_ic ~prefix:str "(...)" then
		[v,Some "(...)"]
	      else
		[]
	    in
	      {str="";
	       id = get_uid ();
	       body =  Some (v ());
	       possibility_list = [];
	       possibilities = 
		  possibilities}
	| S.Record lst ->
	    let v () = 
	      let lst = 
		List.map 
		  (fun (a,b) -> (a,iter (b:Specification.t)))
		  lst
	      in
		Record lst
	    in
	    let possibilities str = 
	      if is_prefix_ic ~prefix:str "{...}" then
		[v,Some "{...}"]
	      else
		[]
	    in
	      {str="";
	       id = get_uid ();
	       body =  Some (v ());
	       possibility_list = [];
	       possibilities = 
		  possibilities}	      

	| S.Atom f -> 
	    let possibilities str : possibility_result = 
	      List.map
		(fun s ->
		   ((fun () -> Atom s),Some s))
		(f str)
	    in
	      {str="";
	       id = get_uid ();
	       body = None;
	       possibility_list = [];
	       possibilities = 
		  possibilities}	      
    in
      fun () -> iter spec
  ;;

  let to_representation : node -> Representation.t = 
    let module R = Representation in
    let rec iter t = 
      match t.body with
	| None ->
	    R.Hole
	| Some x ->
	    iter' x
    and iter' : choices -> Representation.t = 
      function
	| Variant (str,lst) ->
	    R.Variant (str,List.map iter lst)
	| Record lst ->
	    let rec fold = 
	      function
		| [] -> []
		| (str,t) :: xs  ->
		    (str,iter t) :: fold xs
	    in
	      R.Record (fold lst)
	| Tuple lst ->
	    R.Tuple (List.map iter lst)
	| Atom x ->
	    R.Atom x

    in
      iter
  ;;
end
module R = Representation
module S = Specification
module I = Implementation

module Conv = 
struct
  
  let char_of_string str = 
    if String.length str = 1 then 
      String.get str 0 
    else
      failwith "char_of_string"
  ;;

  module Specification = 
  struct
    let to_atom_exn f () = 
      S.Atom
	(fun str -> 
	   try
	     f str;
	     [str]
	   with
	     | _ -> [])
    ;;
    let of_int = to_atom_exn int_of_string
    let of_float = to_atom_exn float_of_string
    let of_bool = to_atom_exn bool_of_string
    let of_char = to_atom_exn char_of_string
    let of_string = to_atom_exn (fun x -> x)
    ;;
    
    let rec of_list a () = 
      S.Variant [("Cons",[a;of_list a]);
		 ("Nil",[])]

    let rec of_unit () = 
      S.Variant ["()",[]]

  end
  module Representation = 
  struct

    
    let of_atom_exn f = 
      function
	| R.Hole -> raise Hole
	| R.Atom s -> 
	    (try
	       f s
	     with
	       | _ -> raise Unexpected)
	| _ -> raise Unexpected
    ;;	
    let to_int = of_atom_exn int_of_string
    let to_bool = of_atom_exn bool_of_string
    let to_float = of_atom_exn float_of_string
    let to_char = of_atom_exn char_of_string
    let to_string = of_atom_exn (fun x -> x)
    ;;
    let to_list _of_a = 
      let rec loop = 
	function
	  | R.Variant ("Cons",[a;b]) ->
	      _of_a a :: loop b
	  | R.Variant ("Nil",[]) ->
	      []
	  | _ -> raise (Of_representation "list")
      in
	loop
    ;;
    let to_unit =
      function
	| R.Variant ("()",[]) ->
	    ()
	| _ -> raise (Of_representation "list")
  end
    
end       


  


module Zipper = 
struct
  module R = Representation
  type t = 
      {parent:t option;
       self:I.node;
       ls: I.node list;
       rs: I.node list
      }
	
  let get_right_sibling 
      {parent=p;self=s;ls=ls;rs=rs}
      =
    match rs with
      | [] -> None
      | x :: xs ->
	  Some 
	    {parent=p;
	     self = x;
	     ls = s :: ls ;
	     rs = xs}

  let get_left_sibling 
      {parent=p;self=s;ls=ls;rs=rs}
      =
    match ls with
      | [] -> None
      | x :: xs ->
	  Some 
	    {parent=p;
	     self = x;
	     ls = xs ;
	     rs = s :: rs}

  let get_child ({parent=p;self=s;ls=ls;rs=rs} as z) : t option = 
    let children = 
      match s.I.body with
	| Some (I.Atom _)
	| None ->
	    []

	| Some (I.Variant (_,lst)) 
	| Some (I.Tuple lst) ->
	    lst
	| Some (I.Record lst) ->
	    List.map snd lst

    in match children with
      | [] -> None
      | x :: xs -> 
	  Some 
	    {self=x;
	     parent=Some z;
	     rs = [];
	     ls = xs
	    }
	    

  let get_parent {parent=p;self=_;ls=_;rs=_} : t option = p
  ;;

end
module Z = Zipper


module Draw = 
struct
  type t = 
      {my:int;
       mx:int;
       vfy:int;
       vly:int;}

  let horizontal_line = Char.code '-' 
  let vertical_line = Char.code '|' 
  let top_corner = Char.code '+' 
  let bottom_corner = Char.code '+' 

  let draw_selector w x dr = 
    if x = dr.mx then 
      ignore (mvwaddstr w x (dr.vly+2) "<-------")
    else
      begin
	for i = dr.vfy + 1 to dr.my + 1 do
	  ignore(mvwaddch w x i horizontal_line)
	done;
	ignore(mvwaddch w x (dr.my+2) top_corner);
	for i = dr.vly + 1 to dr.my + 1 do
	  ignore(mvwaddch w dr.mx i horizontal_line)
	done;
	ignore(mvwaddch w dr.mx (dr.my+2) bottom_corner);
	for i = x + 1 to dr.mx - 1 do
	  ignore(mvwaddch w i (dr.my+2) vertical_line)
	done    

      end
  ;;

  let draw_literal w x y str = 
    ignore(mvwaddstr w x y str);
    let w = y + String.length str in 
      {mx = x;
       my = w;
       vfy=w;
       vly=w}
  ;; 

  let max3 x y z = max (max x y) z
  ;;

  let draw w (currently_selected:I.node) : int * int -> I.node -> t = 
    let rec draw_atom
	?(app="")
	(x,y) 
	({I.id=_;
	  body=body;
	  possibilities=possibilities;
	  str=str;
	  possibility_list=possibility_list} as t)
	=
      let dr = 
	match body with
	  | None
	    ->
	      let pos_str = 
		match possibility_list with
		  | [] -> ""
		  | (str,_) :: xs ->
		      " (" ^ str ^ ")"		      
	      in
	      let str = String.concat "" [str; "__"  ;pos_str] in
		draw_literal w x y str
	  | Some (I.Atom str)
	    ->
	      draw_literal w x y str
	  | Some (I.Variant (str,[])) ->
	      draw_literal w x y str
	  | Some (I.Variant ("Cons",[a;b])) ->
	      let dr = draw_atom ~app:";" (x,y) a in 
	      let dr' = draw_atom (dr.mx+1,y) b in 
		{my = max dr.my dr'.my;
		 mx = dr'.mx;
		 vly = dr'.vly;
		 vfy = dr.vfy}
	  | Some (I.Variant (str,lst))
	    ->
	      ignore(mvwaddstr w x y str);
	      let w = y + String.length str in 
	      let dr = draw_t (x+1,y+2) lst in 
		{my = max w dr.my;
		 mx = dr.mx;
		 vly = dr.vly;
		 vfy = w}	    
	  | Some (I.Tuple lst)
	    ->
	      let rec iter (x,y) : I.node list -> t = 
		function
		  | [t] ->
		      draw_atom (x,y) t 
		  | t :: xs ->
		      let dr = draw_atom ~app:"," (x,y) t in 
		      let dr' = iter (dr.mx+1,y) xs in 
			{my = max dr.my dr'.my;
			 mx = dr'.mx;
			 vly = dr'.vly;
			 vfy = dr.vfy}
		  | [] ->
		      draw_literal w x y "()"
	      in
	      let dr = draw_literal w x y "(" in 
	      let dr = iter (dr.mx,dr.my+1) lst in 
	      let dr' = draw_literal w (dr.mx+1) y ")" in 
		{my = max dr.my dr'.my;
		 mx = dr'.mx;
		 vly = dr'.vly;
		 vfy = dr.vfy}		

	  | Some (I.Record lst) ->
	      let rec iter (x,y) : (string * I.node) list -> t = 
		function
		  | [(str,t)] ->
		      ignore(mvwaddstr w x y (str ^ " =" ));
		      let w = y + String.length str + 2 in
		      let dr = draw_atom (x+1,y+2) t in 
			{my = max w dr.my;
			 mx = dr.mx;
			 vly = dr.vly;
			 vfy = w}
		  | (str,t) :: xs ->
		      ignore(mvwaddstr w x y (str ^ " =" ));
		      let w = y + String.length str + 2 in
		      let dr = draw_atom (x+1,y+2) t in 
		      let dr' = iter (dr.mx+1,y) xs in 
			{my = max3 w dr.my dr'.my;
			 mx = dr'.mx;
			 vly = dr'.vly;
			 vfy = w}
		  | [] ->
		      draw_literal w x y "{}"
	      in
		iter (x,y) lst
      in
      let dr' = draw_literal w dr.mx dr.vly app in 
      let fdr = 
	let vfy = 
	  if dr.mx <> x then 
	    dr.vfy
	  else
	    dr'.vfy
	in
	  {my = max dr.my dr'.my;
	   mx = dr'.mx;
	   vly = dr'.vly;
	   vfy = vfy}	
      in
	if t.I.id == currently_selected.I.id then
	  draw_selector w x fdr;
	fdr
    and draw_t (x,y) t : t = 
      let rec draw_rest (x,y) l = 
	match l with
	  | [] ->
	      failwith "impossible"
	  | [z] ->
	      draw_atom (x,y) z 
	  | z :: zs ->
	      let dr = draw_atom (x,y) z in 
	      let dr' = draw_rest (dr.mx+1,y) zs in
		{my = max dr.my dr'.my;
		 mx = dr'.mx;
		 vly = dr'.vly;
		 vfy = dr.vfy}
      in
	match t with
	  | [] -> 
	      failwith "impossible"
	  | [z] ->
	      draw_atom (x,y) z
	  | z :: zs ->
	      let dr = draw_atom (x,y) z in 
	      let dr' = draw_rest (dr.mx+1,y) zs in 
		{my = max dr.my dr'.my;
		 mx = dr'.mx;
		 vly = dr'.vly;
		 vfy = dr.vfy}
    in
      draw_atom ~app:""

  let draw_zipper w currently_selected (x,y) {Z.parent=p;self=s} = 
    draw w currently_selected  (x,y) s
end

let run ?win spec =     


  ignore (Curses.initscr ());
  ignore (Curses.keypad (Curses.stdscr ()) true);
  ignore (Curses.noecho());
  ignore (Curses.curs_set 0);
  ignore (Curses.refresh ());
  let win = 
    match win with
      | Some w -> w
      | None -> Curses.stdscr ()
  in
  let test_t = I.of_specification spec () in
  let currently_selected : Z.t ref = ref {Z.parent=None;self=test_t;ls=[];rs=[]} in 
  let draw_all () = 
    ignore(Curses.wclear win);
    ignore(Draw.draw win !currently_selected.Z.self (0,0) test_t);
    ignore(Curses.wrefresh win);
  in
  let rotate_possibility_list t = 
    match t.I.possibility_list with
      | [] -> ()
      | x :: xs ->
	  t.I.possibility_list <- xs @ [x]
  in
  let update_possibility_list t = 
    let possibility_list = t.I.possibilities t.I.str in 
    let f (c,str_opt) =
      match str_opt with
	| None ->
	    (t.I.str,c)
	| Some s ->
	    (s,c)
    in
    let possibility_list = List.map f possibility_list in 
      t.I.possibility_list <- possibility_list
  in
  let update_string t str = 
    t.I.str <- str;
    update_possibility_list t
  in
	
	
    



    
    let rec loop () =
      draw_all ();
      let ch = Curses.getch () in 
	if ch = Key.left then 
	  match 
	    Z.get_parent !currently_selected
	  with
	    | None ->
		()
	    | Some x ->
		currently_selected := x
	else  if ch = Key.right then 
 	  match 
	    Z.get_child !currently_selected
	  with
	    | None ->
		()
	    | Some x ->
		currently_selected := x
	else if ch = Key.down then 
 	  match 
	    Z.get_left_sibling !currently_selected
	  with
	    | None ->
		()
	    | Some x ->
		currently_selected := x
	else if ch = Key.up then 
	  begin
 	    match 
	      Z.get_right_sibling !currently_selected
	    with
	      | None ->
		  ()
	      | Some x ->
		  currently_selected := x
	  end
	else if ch = Key.dc then
	  begin
	    match !currently_selected.Z.self with
	      | {I.str=str;
		 id=_;
		 body = Some _ ;}
		->
		  let t = !currently_selected.Z.self in
		  t.I.body <- None;
		  update_possibility_list t
	      | _ -> ()
	  end
	else if ch = (*Key.enter*) 10 then (* CR: Key.enter does not work *)
	  begin
	    match !currently_selected.Z.self with
	      | ({I.str=str;
		  id=_;
		  possibilities = _;
		  possibility_list = possibility_list;
		  body = None} as z)
		->	
		  begin 
		    match possibility_list with
		      | [] -> ()
		      |  (_,t') :: _ ->
			   z.I.body <- Some (t' ());
			   update_possibility_list z
		  end
	      | _ -> ()
	  end
	else if ch = Key.backspace then
	  begin
	    match !currently_selected.Z.self with
	      | ({I.str=str;
		  id=_;
		  body = None} as z)
		->	
		  let len = String.length str in
		    if len > 0 then
			update_string z (String.sub str 0 (len - 1));

	      | _ -> ()

	  end
	else if ch = 9 (* where's the backspace key? *) then
	  begin
	    let t = !currently_selected.Z.self in 
	      match t.I.body with
		| None -> rotate_possibility_list t
		| Some _ -> ()
	  end
	else
	  begin
	    match !currently_selected.Z.self with
	      | ({I.str=str;
		  id=_;
		  body = None} as z)
		->	
		  if Char.chr ch = '\t' then 
		    update_string z (str ^ string_of_int ch)
		  else
		  (try
		     update_string z (str ^ (String.make 1 (Char.chr ch)))
		   with
		     | _ -> update_string z (str ^ ("\\" ^ string_of_int ch)))
	      | _ -> ()
	  end;

	loop ()
    in
      loop ()
;;
